import os
import time

import requests
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options
from webdriver_manager.chrome import ChromeDriverManager


class BilibiliDanmakuDownloader:
    def __init__(self, video_url):
        self.video_url = video_url
        self.bvid = self.extract_bvid()

        # 设置Chrome选项
        self.chrome_options = Options()
        self.chrome_options.add_argument('--headless')  # 无头模式
        self.chrome_options.add_argument('--disable-gpu')
        self.chrome_options.add_argument('--no-sandbox')

        # 初始化WebDriver
        self.driver = webdriver.Chrome(
            service=Service("D:\Program Files\Python\python\chromedriver.exe"),
            options=self.chrome_options
        )

    def extract_bvid(self):
        """从URL中提取BV号"""
        if 'BV' in self.video_url:
            return self.video_url.split('BV')[1][:12]
        return None

    def get_cid(self):
        """获取视频CID"""
        self.driver.get(self.video_url)
        time.sleep(3)

        # 从页面源码中提取CID
        cid = self.driver.execute_script("""
            return window.__INITIAL_STATE__?.videoData?.cid || 
                   window.__INITIAL_STATE__?.epInfo?.cid;
        """)
        return cid

    def download_danmaku(self, segment_index=1):
        """下载指定分段的弹幕"""
        cid = self.get_cid()
        if not cid:
            print("无法获取视频CID")
            return False

        # 构造弹幕API URL
        danmaku_url = f"https://api.bilibili.com/x/v2/dm/web/seg.so?type=1&oid={cid}&segment_index={segment_index}"

        # 设置请求头
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/137.0.0.0 Safari/537.36 Edg/137.0.0.0',
            'Referer': self.video_url
        }

        try:
            response = requests.get(danmaku_url, headers=headers)
            if response.status_code == 200:
                # 保存为seg.so文件
                with open('弹幕/seg.so', 'wb') as f:
                    f.write(response.content)
                print(f"成功下载弹幕文件: seg.so (CID: {cid})")
                return True
            else:
                print(f"下载失败，状态码: {response.status_code}")
        except Exception as e:
            print(f"下载出错: {e}")
        return False

    def close(self):
        """关闭浏览器"""
        self.driver.quit()


# 使用示例
if __name__ == "__main__":
    # 替换为你要抓取的B站视频URL
    video_url = "https://www.bilibili.com/video/BV1YhNqzsE65/?spm_id_from=333.1007.tianma.1-1-1.click&vd_source=6e299066bd0191296af111801fbae94a"

    downloader = BilibiliDanmakuDownloader(video_url)
    if downloader.download_danmaku(segment_index=1):  # 下载第1段弹幕
        print("弹幕文件已保存为 seg.so")
    downloader.close()